import numpy as np
import cv2
import matplotlib.pyplot as plt
from PIL import Image
from bm3d import bm3d, BM3DProfile
from skimage import util


# 设置全局字体为支持中文的字体(如 SimHei、Microsoft YaHei)
plt.rcParams['font.sans-serif'] = ['SimHei']  # Windows 系统常用字体
# plt.rcParams['font.sans-serif'] = ['Arial Unicode MS']  # macOS 系统常用字体
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

# 读取RGB图像
image_rgb =  np.array(Image.open('frame_18008_tex_in.png').convert("RGB"))

# 将图像转换为浮点数格式并归一化到[0, 1]范围
image_float = image_rgb.astype(np.float32) / 255.0

# 分离R、G、B三个通道
r_channel = image_float[:, :, 0]
g_channel = image_float[:, :, 1]
b_channel = image_float[:, :, 2]

# 添加高斯噪声（可选）
noise_sigma = 0.05  # 噪声标准差
# noisy_r = util.random_noise(r_channel, mode='gaussian', var=noise_sigma**2)
# noisy_g = util.random_noise(g_channel, mode='gaussian', var=noise_sigma**2)
# noisy_b = util.random_noise(b_channel, mode='gaussian', var=noise_sigma**2)

# 对每个通道分别使用BM3D进行降噪
# 基本版本（速度较快）
denoised_r_basic = bm3d(r_channel, noise_sigma)
denoised_g_basic = bm3d(g_channel, noise_sigma)
denoised_b_basic = bm3d(b_channel, noise_sigma)

# 高级版本（质量更好但速度较慢）
profile = BM3DProfile()  # 使用默认配置
denoised_r = bm3d(r_channel, noise_sigma, profile=profile)
denoised_g = bm3d(g_channel, noise_sigma, profile=profile)
denoised_b = bm3d(b_channel, noise_sigma, profile=profile)

# 合并三个通道的结果
denoised_image_basic = np.stack([denoised_r_basic, denoised_g_basic, denoised_b_basic], axis=2)
denoised_image = np.stack([denoised_r, denoised_g, denoised_b], axis=2)

# 将降噪后的图像转换回0-255范围并转换为uint8格式
denoised_image_basic_uint8 = (denoised_image_basic * 255).astype(np.uint8)
denoised_image_uint8 = (denoised_image * 255).astype(np.uint8)

# 显示结果
plt.figure(figsize=(15, 10))

plt.subplot(2, 2, 1)
plt.imshow(image_rgb)
plt.title('原始RGB图像')
plt.axis('off')

plt.subplot(2, 2, 2)
plt.imshow(image_rgb)
plt.title(f'原始RGB图像 (噪声σ={noise_sigma})')
plt.axis('off')

plt.subplot(2, 2, 3)
plt.imshow(denoised_image_basic)
plt.title('BM3D基本版本RGB降噪结果')
plt.axis('off')

plt.subplot(2, 2, 4)
plt.imshow(denoised_image)
plt.title('BM3D高级版本RGB降噪结果')
plt.axis('off')

plt.tight_layout()
plt.show()

# save denoise image
Image.fromarray(denoised_image_basic_uint8).save('denoised_image_basic.png')
Image.fromarray(denoised_image_uint8).save('denoised_image.png')
